import logging
import json
import os
import shutil
import socket
import tempfile
from mcp.embedding import Embedding
from fastapi import FastAPI, HTTPException, File, UploadFile, Form, Request
from fastapi.responses import JSONResponse, FileResponse
from pydantic import BaseModel
import uvicorn

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger('mcp.server')

# 初始化FastAPI应用
app = FastAPI(title="MCP Embedding Service", version="1.0")

# 加载配置
def load_config(config_path):
    """加载配置文件"""
    if os.path.exists(config_path):
        try:
            with open(config_path, 'r', encoding='utf-8') as f:
                return json.load(f)
        except Exception as e:
            logger.error(f"加载配置文件失败: {str(e)}")
            return {}
    else:
        logger.warning(f"配置文件 {config_path} 不存在，使用默认配置")
        return {}

# 定义请求模型
class ProcessDocumentRequest(BaseModel):
    content: str = ''
    
    def validate_content(self):
        if not self.content.strip():
            raise HTTPException(status_code=400, detail="内容不能为空")
        return self

class QueryDocumentsRequest(BaseModel):
    query: str
    top_k: int = 3

# 定义API路由
@app.post("/api/v1/process_document")
async def process_document(request: Request):
    """处理文档并保存嵌入"""
    try:
        # 处理JSON请求
        request_data = await request.json()
        content = request_data.get('content', '')
        
        # 检查内容是否为空
        if not content.strip():
            raise HTTPException(status_code=400, detail="内容不能为空")
        
        # 检查是否为base64编码的内容
        is_base64 = False
        try:
            import base64
            # 尝试解码base64内容
            decoded_content = base64.b64decode(content)
            is_base64 = True
        except Exception:
            # 不是base64编码，使用原始内容
            decoded_content = content
        
        if is_base64:
            logger.info(f"接收到base64编码内容，解码后大小: {len(decoded_content)} 字节")
            # 对于二进制内容，我们可能需要保存为临时文件并处理
            with tempfile.NamedTemporaryFile(delete=False) as temp_file:
                temp_file.write(decoded_content)
                temp_file_path = temp_file.name
            
            try:
                result = embedding_service.embedding(file_path=temp_file_path)
            finally:
                # 确保临时文件被删除
                os.unlink(temp_file_path)
            logger.info(f"接收到Base64编码文本内容，长度: {len(decoded_content)} 字符")
        else:
            logger.info(f"接收到文本内容，长度: {len(content)} 字符")
            result = embedding_service.embedding(content=content)
        
        # 确保返回值是有效的JSON对象
        if result is None:
            return {'status': 'success', 'message': '文档处理成功'}        
        return result
    except HTTPException as e:
        logger.warning(f"文档验证失败: {str(e.detail)}")
        raise e
    except Exception as e:
        logger.error(f"处理文档时出错: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

# 处理图片和描述文本接口
@app.post("/api/v1/process_image_caption")
def process_image_caption(
    caption: str = Form(...),
    image: UploadFile = File(...)
):
    try:
        # 保存图片到临时路径
        temp_image_filename = f"temp_{os.urandom(8).hex()}{os.path.splitext(image.filename)[1]}"
        temp_image_path = os.path.join(IMAGE_STORAGE_DIR, temp_image_filename)
        with open(temp_image_path, "wb") as buffer:
            shutil.copyfileobj(image.file, buffer)

        # 获取当前图片的大小
        current_image_size = os.path.getsize(temp_image_path)
        logger.info(f"当前图片大小: {current_image_size} 字节")

        # 检查是否存在相同描述和相同内容的图片
        # 1. 先进行精确描述匹配（去除前后空白字符）
        cleaned_caption = caption.strip()
        exact_match = embedding_service.query_exact_document(document=cleaned_caption)
        logger.info(f"精确匹配结果: {exact_match}")
        
        if exact_match.get('status') == 'success' and 'results' in exact_match and len(exact_match['results']) > 0:
            for result in exact_match['results']:
                if 'image_path' in result and result['image_path'] and 'image_size' in result:
                    # 检查图片文件是否存在
                    if os.path.exists(result['image_path']):
                        # 比较图片大小（作为初步检查）
                        if result['image_size'] == current_image_size:
                            # 计算图片哈希值进行精确比较
                            def calculate_image_hash(image_path):
                                import hashlib
                                hasher = hashlib.md5()
                                with open(image_path, 'rb') as f:
                                    buf = f.read()
                                    hasher.update(buf)
                                return hasher.hexdigest()

                            existing_image_hash = calculate_image_hash(result['image_path'])
                            current_image_hash = calculate_image_hash(temp_image_path)

                            if existing_image_hash == current_image_hash:
                                logger.info(f"已存在相同描述和内容的图片: {result['image_path']}, 描述: {caption}")
                                # 删除临时图片
                                os.remove(temp_image_path)
                                return JSONResponse(content={"status": "success", "image_path": result['image_path'], "message": "图片已存在，跳过保存"}, status_code=200)
                            else:
                                logger.info(f"存在相同描述但内容不同的图片: {result['image_path']}")
                        else:
                            logger.info(f"存在相同描述但大小不同的图片: {result['image_path']}")

        # 2. 如果没有精确匹配，进行相似描述查询（作为备选）
        query_result = embedding_service.query(query=caption, top_k=10)
        logger.info(f"相似查询结果: {query_result}")

        # 重命名临时图片为正式名称
        image_filename = f"{os.path.splitext(image.filename)[0]}_{os.urandom(8).hex()}{os.path.splitext(image.filename)[1]}"
        image_path = os.path.join(IMAGE_STORAGE_DIR, image_filename)
        os.rename(temp_image_path, image_path)
        logger.info(f"保存图片到: {image_path}")

        # 处理描述文本
        embedding_service.process_document(content=caption, is_image_caption=True, image_path=image_path, image_size=current_image_size)
        return JSONResponse(content={"status": "success", "image_path": image_path}, status_code=200)
    except Exception as e:
        logger.error(f"处理图片和描述文本时出错: {str(e)}")
        raise HTTPException(status_code=500, detail=f"处理图片和描述文本时出错: {str(e)}")

@app.post("/api/v1/query")
def query(request: QueryDocumentsRequest):
    """查询相关文档"""
    try:
        result = embedding_service.query(
            query=request.query,
            top_k=request.top_k
        )
        logger.info(f"查询结果: {result}")
        print("")

        # 处理结果，添加图片URL
        if isinstance(result, list):
            for item in result:
                if isinstance(item, dict) and 'image_path' in item and item['image_path']:
                    # 提取图片文件名
                    image_filename = os.path.basename(item['image_path'])
                    # 构建图片URL，包含协议、主机和端口
                    item['image_url'] = f"http://localhost:{port}/api/v1/image/{image_filename}"
                    logger.info(f"添加图片URL: {item['image_url']}")
                else:
                    item['image_url'] = ""
        
        return result
    except Exception as e:
        logger.error(f"查询文档时出错: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

# 获取图片接口
@app.get("/api/v1/image/{image_filename}")
def get_image(image_filename: str):
    image_path = os.path.join(IMAGE_STORAGE_DIR, image_filename)
    logger.info(f"尝试获取图片: {image_path}")
    if not os.path.exists(image_path):
        logger.error(f"图片不存在: {image_path}")
        raise HTTPException(status_code=404, detail=f"图片不存在: {image_path}")
    # 检查文件权限
    if not os.access(image_path, os.R_OK):
        logger.error(f"没有权限读取图片: {image_path}")
        raise HTTPException(status_code=403, detail=f"没有权限读取图片: {image_path}")
    logger.info(f"成功获取图片: {image_path}")
    return FileResponse(image_path)

# 健康检查
@app.get("/health")
def health_check():
    return {"status": "healthy"}

@app.get("/")
def health_check():
    return "Welcome to RAG MCP Server"

def get_all_ip_addresses() -> list[str]:
    """
    返回本机所有 IPv4 地址（含 127.0.0.1）
    """
    ips = set()

    # 1. 通用方法：socket.gethostbyname_ex
    try:
        _, _, addrs = socket.gethostbyname_ex(socket.gethostname())
        ips.update(addrs)
    except socket.gaierror:
        pass
    
    return ips

# 启动服务
if __name__ == "__main__":
    logger.info("启动MCP服务...")

    # 加载配置
    config = load_config(config_path='../config/config.json')
    # 从配置中提取Embedding类所需的参数
    embedding_config = {
        'embedding_model_name': config.get('embedding_model', 'shibing624/text2vec-base-chinese'),
        'embedding_model_path': config.get('embedding_model_path'),
        'rerank_model_name': config.get('rerank_model', 'cross-encoder/mmarco-mMiniLMv2-L12-H384-v1'),
        'rerank_model_path': config.get('rerank_model_path'),
        'collection_name': config.get('collection_name', 'rag_collection')
    }
    embedding_service = Embedding(**embedding_config)

    # 创建图片存储目录
    IMAGE_STORAGE_DIR = os.path.abspath(config.get('image_storage_path', '../image_storage'))
    os.makedirs(IMAGE_STORAGE_DIR, exist_ok=True)
    logger.info(f"图片存储目录: {IMAGE_STORAGE_DIR}")

    # 从配置文件获取主机和端口
    host = config.get('server_host', '0.0.0.0')
    port = int(config.get('server_port', 8000))

    actual_ip_list = get_all_ip_addresses() if host == "0.0.0.0" else host
    logger.info(f"服务将运行在 http://{host}:{port}")
    for actual_ip in actual_ip_list:
         logger.info(f"可通过URL访问: http://{actual_ip}:{port}")
    uvicorn.run(app, host=host, port=port)